import torch
from torch.utils.tensorboard import SummaryWriter
import os
from datetime import datetime


class TensorBoardLogger:
    """TensorBoard日志记录器"""
    
    def __init__(self, log_dir, model_name, comment=""):
        """
        初始化TensorBoard记录器
        
        Args:
            log_dir (str): 日志目录
            model_name (str): 模型名称
            comment (str): 附加注释
        """
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.log_path = os.path.join(log_dir, f'{model_name}_{timestamp}_{comment}')
        os.makedirs(self.log_path, exist_ok=True)
        
        self.writer = SummaryWriter(self.log_path)
        print(f'TensorBoard日志保存至: {self.log_path}')
    
    def log_scalars(self, tag_scalar_dict, global_step):
        """
        记录标量值
        
        Args:
            tag_scalar_dict (dict): 标签-数值字典
            global_step (int): 全局步数
        """
        for tag, scalar in tag_scalar_dict.items():
            self.writer.add_scalar(tag, scalar, global_step)
    
    def log_scalar(self, tag, scalar, global_step):
        """
        记录单个标量值
        
        Args:
            tag (str): 标签
            scalar (float): 数值
            global_step (int): 全局步数
        """
        self.writer.add_scalar(tag, scalar, global_step)
    
    def log_scalars_dict(self, main_tag, tag_scalar_dict, global_step):
        """
        记录标量字典（用于对比图）
        
        Args:
            main_tag (str): 主标签
            tag_scalar_dict (dict): 标签-数值字典
            global_step (int): 全局步数
        """
        self.writer.add_scalars(main_tag, tag_scalar_dict, global_step)
    
    def log_model_graph(self, model, input_tensor):
        """
        记录模型结构图
        
        Args:
            model (torch.nn.Module): 模型
            input_tensor (torch.Tensor): 示例输入
        """
        try:
            self.writer.add_graph(model, input_tensor)
            print('模型结构已记录到TensorBoard')
        except Exception as e:
            print(f'记录模型结构时出错: {e}')
    
    def log_text(self, tag, text_string):
        """
        记录文本信息
        
        Args:
            tag (str): 标签
            text_string (str): 文本内容
        """
        self.writer.add_text(tag, text_string)
    
    def log_training_epoch(self, epoch, train_loss, train_acc, val_loss, val_acc, lr):
        """
        记录训练epoch的指标
        
        Args:
            epoch (int): epoch数
            train_loss (float): 训练损失
            train_acc (float): 训练准确率
            val_loss (float): 验证损失
            val_acc (float): 验证准确率
            lr (float): 学习率
        """
        # 记录损失和准确率
        self.writer.add_scalar('Loss/Train', train_loss, epoch)
        self.writer.add_scalar('Loss/Validation', val_loss, epoch)
        self.writer.add_scalar('Accuracy/Train', train_acc, epoch)
        self.writer.add_scalar('Accuracy/Validation', val_acc, epoch)
        self.writer.add_scalar('Learning_Rate', lr, epoch)
        
        # 记录对比图
        self.writer.add_scalars('Loss_Comparison', 
                               {'Train': train_loss, 'Validation': val_loss}, epoch)
        self.writer.add_scalars('Accuracy_Comparison', 
                               {'Train': train_acc, 'Validation': val_acc}, epoch)
    
    def log_training_batch(self, global_step, loss, accuracy):
        """
        记录训练batch的指标
        
        Args:
            global_step (int): 全局步数
            loss (float): 批次损失
            accuracy (float): 批次准确率
        """
        self.writer.add_scalar('Train/Batch_Loss', loss, global_step)
        self.writer.add_scalar('Train/Batch_Accuracy', accuracy, global_step)
    
    def log_best_accuracy(self, best_acc, epoch):
        """
        记录最佳准确率
        
        Args:
            best_acc (float): 最佳准确率
            epoch (int): epoch数
        """
        self.writer.add_scalar('Best_Validation_Accuracy', best_acc, epoch)
    
    def log_test_results(self, test_acc):
        """
        记录测试结果
        
        Args:
            test_acc (float): 测试准确率
        """
        self.writer.add_scalar('Test/Final_Accuracy', test_acc, 0)
    
    def log_training_summary(self, training_time, best_val_acc, num_epochs, batch_size, lr):
        """
        记录训练总结
        
        Args:
            training_time (float): 训练时间
            best_val_acc (float): 最佳验证准确率
            num_epochs (int): 总epoch数
            batch_size (int): 批次大小
            lr (float): 初始学习率
        """
        summary = f'''
**训练总结**
- 训练时间: {training_time:.2f} 秒
- 最佳验证准确率: {best_val_acc:.2f}%
- 总epoch数: {num_epochs}
- 批次大小: {batch_size}
- 初始学习率: {lr}
'''
        self.writer.add_text('Training_Summary', summary)
    
    def log_final_summary(self, training_time, best_val_acc, test_acc, model_path):
        """
        记录最终结果总结
        
        Args:
            training_time (float): 训练时间
            best_val_acc (float): 最佳验证准确率
            test_acc (float): 测试准确率
            model_path (str): 模型保存路径
        """
        summary = f'''
**最终结果总结**
- 训练时间: {training_time:.2f} 秒
- 最佳验证准确率: {best_val_acc:.2f}%
- 最终测试准确率: {test_acc:.2f}%
- 模型保存路径: {model_path}
- TensorBoard日志路径: {self.log_path}
'''
        self.writer.add_text('Final_Summary', summary)
    
    def close(self):
        """关闭writer"""
        self.writer.close()
    
    def get_log_path(self):
        """获取日志路径"""
        return self.log_path
    
    def print_tensorboard_command(self):
        """打印TensorBoard启动命令"""
        print(f'TensorBoard日志已保存，可使用以下命令查看:')
        print(f'tensorboard --logdir={self.log_path}')
        print('可在浏览器中访问 http://localhost:6006 查看训练过程可视化') 